import torch
import numpy as np
import torch.nn as nn
from timm.models.layers import trunc_normal_
from einops import rearrange, repeat
from torch_scatter import scatter_softmax, scatter
from utils.normalization import Normalizer

ACTIVATION = {
    "gelu": nn.GELU,
    "tanh": nn.Tanh,
    "sigmoid": nn.Sigmoid,
    "relu": nn.ReLU,
    "leaky_relu": nn.LeakyReLU(0.1),
    "softplus": nn.Softplus,
    "ELU": nn.ELU,
    "silu": nn.SiLU,
}


def count_model_parameters(model: nn.Module) -> int:
    """
    统计PyTorch模型的参数量。

    参数:
    model (nn.Module): 需要统计参数量的PyTorch模型。

    返回值:
    int: 模型的总参数量。
    """
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params


class MLP(nn.Module):
    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act="gelu", res=True):
        super(MLP, self).__init__()

        if act in ACTIVATION.keys():
            act = ACTIVATION[act]
        else:
            raise NotImplementedError
        self.n_input = n_input
        self.n_hidden = n_hidden
        self.n_output = n_output
        self.n_layers = n_layers
        self.res = res
        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act())
        self.linear_post = nn.Linear(n_hidden, n_output)
        self.linears = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(n_hidden, n_hidden), act())
                for _ in range(n_layers)
            ]
        )

    def forward(self, x):
        x = self.linear_pre(x)
        for i in range(self.n_layers):
            if self.res:
                x = self.linears[i](x) + x
            else:
                x = self.linears[i](x)
        x = self.linear_post(x)
        return x


class Physics_Attention_1D(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, slice_num=64):
        super().__init__()
        inner_dim = dim_head * heads
        self.dim_head = dim_head
        self.heads = heads
        self.scale = dim_head**-0.5
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)

        self.in_project_x = nn.Linear(dim, inner_dim)
        self.in_project_fx = nn.Linear(dim, inner_dim)
        self.in_project_slice = nn.Linear(dim_head, slice_num)
        for l in [self.in_project_slice]:
            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def forward(self, x):
        # B N C
        B, N, C = x.shape

        ### (1) Slice
        fx_mid = (
            self.in_project_fx(x)
            .reshape(B, N, self.heads, self.dim_head)
            .permute(0, 2, 1, 3)
            .contiguous()
        )  # B H N C
        x_mid = (
            self.in_project_x(x)
            .reshape(B, N, self.heads, self.dim_head)
            .permute(0, 2, 1, 3)
            .contiguous()
        )  # B H N C
        slice_weights = self.softmax(
            self.in_project_slice(x_mid) / self.temperature
        )  # B H N G
        slice_norm = slice_weights.sum(2)  # B H G
        slice_token = torch.einsum("bhnc,bhng->bhgc", fx_mid, slice_weights)
        # torch.einsum("bhcn,bhng->bhgc", fx_mid, slice_weights)
        slice_token = slice_token / (
            (slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head)
        )

        ### (2) Attention among slice tokens
        q_slice_token = self.to_q(slice_token)
        k_slice_token = self.to_k(slice_token)
        v_slice_token = self.to_v(slice_token)
        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
        attn = self.softmax(dots)
        attn = self.dropout(attn)
        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D

        ### (3) Deslice
        out_x = torch.einsum("bhgc,bhng->bhnc", out_slice_token, slice_weights)
        out_x = rearrange(out_x, "b h n d -> b n (h d)")
        return self.to_out(out_x)


class GraphPhysics_Attention_1D(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, slice_num=64):
        super().__init__()
        inner_dim = dim_head * heads
        self.dim_head = dim_head
        self.heads = heads
        self.scale = dim_head**-0.5
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)

        self.in_project_x = nn.Linear(dim, inner_dim)
        self.in_project_fx = nn.Linear(dim, inner_dim)
        self.in_project_slice = nn.Linear(dim_head, slice_num)
        for l in [self.in_project_slice]:
            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def forward(self, x, graph_node):

        # B N C
        BN, C = x.shape

        ### (1) Slice
        fx_mid = (
            self.in_project_fx(x)
            .reshape(BN, self.heads, self.dim_head)
            .permute(1, 0, 2)
            .contiguous()
        )  # H BN  C

        x_mid = (
            self.in_project_x(x)
            .reshape(BN, self.heads, self.dim_head)
            .permute(1, 0, 2)
            .contiguous()
        )  #  H BN C

        slice_weights = scatter_softmax(
            src=self.in_project_slice(x_mid) / self.temperature,
            index=graph_node.batch,
            dim=-1,
        )
        slice_norm = scatter(
            slice_weights, graph_node.batch, dim=1, reduce="sum"
        )  # B H G
        slice_token = torch.einsum("bhnc,bhng->bhgc", fx_mid, slice_weights)
        # torch.einsum("bhcn,bhng->bhgc", fx_mid, slice_weights)
        slice_token = slice_token / (
            (slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head)
        )

        ### (2) Attention among slice tokens
        q_slice_token = self.to_q(slice_token)
        k_slice_token = self.to_k(slice_token)
        v_slice_token = self.to_v(slice_token)
        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
        attn = self.softmax(dots)
        attn = self.dropout(attn)
        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D

        ### (3) Deslice
        out_x = torch.einsum("bhgc,bhng->bhnc", out_slice_token, slice_weights)
        out_x = rearrange(out_x, "b h n d -> b n (h d)")
        return self.to_out(out_x)


class GraphTransolver_block(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dropout: float,
        act="gelu",
        mlp_ratio=4,
        last_layer=False,
        out_dim=1,
        slice_num=32,
    ):
        super().__init__()
        self.last_layer = last_layer
        self.ln_1 = nn.LayerNorm(hidden_dim)
        self.Attn = GraphPhysics_Attention_1D(
            hidden_dim,
            heads=num_heads,
            dim_head=hidden_dim // num_heads,
            dropout=dropout,
            slice_num=slice_num,
        )
        self.ln_2 = nn.LayerNorm(hidden_dim)
        self.mlp = MLP(
            hidden_dim,
            hidden_dim * mlp_ratio,
            hidden_dim,
            n_layers=0,
            res=False,
            act=act,
        )
        if self.last_layer:
            self.ln_3 = nn.LayerNorm(hidden_dim)
            self.mlp2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, fx, graph_node):
        fx = self.Attn(self.ln_1(fx), graph_node) + fx
        fx = self.mlp(self.ln_2(fx)) + fx
        if self.last_layer:
            return self.mlp2(self.ln_3(fx))
        else:
            return fx


class Transolver_block(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dropout: float,
        act="gelu",
        mlp_ratio=4,
        last_layer=False,
        out_dim=1,
        slice_num=32,
    ):
        super().__init__()
        self.last_layer = last_layer
        self.ln_1 = nn.LayerNorm(hidden_dim)
        self.Attn = Physics_Attention_1D(
            hidden_dim,
            heads=num_heads,
            dim_head=hidden_dim // num_heads,
            dropout=dropout,
            slice_num=slice_num,
        )
        self.ln_2 = nn.LayerNorm(hidden_dim)
        self.mlp = MLP(
            hidden_dim,
            hidden_dim * mlp_ratio,
            hidden_dim,
            n_layers=0,
            res=False,
            act=act,
        )
        if self.last_layer:
            self.ln_3 = nn.LayerNorm(hidden_dim)
            self.mlp2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, fx):
        fx = self.Attn(self.ln_1(fx)) + fx
        fx = self.mlp(self.ln_2(fx)) + fx
        if self.last_layer:
            return self.mlp2(self.ln_3(fx))
        else:
            return fx


class Multiscale_Transolver_block(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dropout: float,
        act="gelu",
        mlp_ratio=4,
        last_layer=False,
        out_dim=1,
        slice_nums=[8, 16, 32, 64],
    ):
        super().__init__()
        self.last_layer = last_layer
        self.ln_1 = nn.LayerNorm(hidden_dim)
        self.Attns = nn.ModuleList(
            [
                Physics_Attention_1D(
                    hidden_dim,
                    heads=num_heads,
                    dim_head=hidden_dim // num_heads,
                    dropout=dropout,
                    slice_num=slice_num,
                )
                for slice_num in slice_nums
            ]
        )
        self.ln_2s = nn.ModuleList(
            [nn.LayerNorm(hidden_dim) for slice_num in slice_nums]
        )
        self.mlp = MLP(
            hidden_dim * len(slice_nums),
            hidden_dim * mlp_ratio,
            hidden_dim,
            n_layers=0,
            res=False,
            act=act,
        )
        if self.last_layer:
            self.ln_3 = nn.LayerNorm(hidden_dim)
            self.mlp2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, fx):
        norm_fx = self.ln_1(fx)
        multi_scale_atts = [
            ln_2(att(norm_fx) + fx) for att, ln_2 in zip(self.Attns, self.ln_2s)
        ]
        multi_scale_atts = torch.cat(multi_scale_atts, dim=-1)
        fx = self.mlp(multi_scale_atts) + fx
        if self.last_layer:
            return self.mlp2(self.ln_3(fx))
        else:
            return fx


# mt = Multiscale_Transolver_block(hidden_dim=32, dropout=0, num_heads=8)
# rst = mt(torch.randn([1, 3148, 32]))
# print(rst.shape)


class Model(nn.Module):
    def __init__(
        self,
        space_dim=1,
        n_layers=5,
        n_hidden=256,
        dropout=0,
        n_head=8,
        act="gelu",
        mlp_ratio=1,
        fun_dim=1,
        out_dim=1,
        slice_num=32,
        ref=8,
        unified_pos=False,
    ):
        super(Model, self).__init__()
        self.__name__ = "UniPDE_3D"
        self.ref = ref
        self.unified_pos = unified_pos
        if self.unified_pos:
            self.preprocess = MLP(
                fun_dim + self.ref * self.ref * self.ref,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )
        else:
            self.preprocess = MLP(
                fun_dim + space_dim,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )

        self.n_hidden = n_hidden
        self.space_dim = space_dim

        self.blocks = nn.ModuleList(
            [
                Transolver_block(
                    num_heads=n_head,
                    hidden_dim=n_hidden,
                    dropout=dropout,
                    act=act,
                    mlp_ratio=mlp_ratio,
                    out_dim=out_dim,
                    slice_num=slice_num,
                    last_layer=False,
                )
                for _ in range(n_layers)
            ]
        )

        self.decoder = nn.Sequential(
            nn.LayerNorm(n_hidden), nn.Linear(n_hidden, out_dim)
        )

        self.initialize_weights()
        self.placeholder = nn.Parameter(
            (1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float)
        )

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_grid(self, my_pos):
        # my_pos 1 N 3
        batchsize = my_pos.shape[0]

        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat(
            [batchsize, 1, self.ref, self.ref, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)
        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat(
            [batchsize, self.ref, 1, self.ref, 1]
        )
        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat(
            [batchsize, self.ref, self.ref, 1, 1]
        )
        grid_ref = (
            torch.cat((gridx, gridy, gridz), dim=-1)
            .cuda()
            .reshape(batchsize, self.ref**3, 3)
        )  # B 4 4 4 3

        pos = (
            torch.sqrt(
                torch.sum(
                    (my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2, dim=-1
                )
            )
            .reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref)
            .contiguous()
        )
        return pos

    def encode(self, x, pos):

        # cfd_data, geom_data = data
        # x, fx, T = cfd_data.x, None, None
        # x  : node features
        # pos: node pos

        if self.unified_pos:
            new_pos = self.get_grid(pos[None, :, :])
            x = torch.cat((x, new_pos), dim=-1)

        pos, fx = None, None

        if fx is not None:
            fx = torch.cat((x, fx), -1)
            fx = self.preprocess(fx)
        else:
            fx = self.preprocess(x)
            fx = fx + self.placeholder[None, None, :]

        return fx

    def process(self, fx):

        for block in self.blocks:
            fx = block(fx)

        return fx

    def decode(self, fx, num_graphs=1):

        fx = self.decoder(fx)

        if len(fx.shape) == 3:
            fx = fx.reshape(num_graphs * 3586, -1)

        return fx

    def forward(self, x, graph_node, params=None):
        # cfd_data, geom_data = data
        # x, fx, T = cfd_data.x, None, None
        # x  : node features
        # pos: node pos

        if len(x.shape) == 2:
            x = x.reshape(graph_node.num_graphs, 3586, -1)

        fx = self.encode(x, pos=graph_node.pos)

        fx = self.process(fx)

        fx = self.decode(fx, graph_node.num_graphs)

        return fx


class MS_Model(nn.Module):
    def __init__(
        self,
        space_dim=1,
        n_layers=5,
        n_hidden=256,
        dropout=0,
        n_head=8,
        act="gelu",
        mlp_ratio=1,
        fun_dim=1,
        out_dim=1,
        slice_nums=[8, 16, 32, 64],
        ref=8,
        unified_pos=False,
    ):
        super(MS_Model, self).__init__()
        self.__name__ = "UniPDE_3D"
        self.ref = ref
        self.unified_pos = unified_pos
        if self.unified_pos:
            self.preprocess = MLP(
                fun_dim + self.ref * self.ref * self.ref,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )
        else:
            self.preprocess = MLP(
                fun_dim + space_dim,
                n_hidden * 2,
                n_hidden,
                n_layers=0,
                res=False,
                act=act,
            )

        self.n_hidden = n_hidden
        self.space_dim = space_dim

        self.blocks = nn.ModuleList(
            [
                Multiscale_Transolver_block(
                    num_heads=n_head,
                    hidden_dim=n_hidden,
                    dropout=dropout,
                    act=act,
                    mlp_ratio=mlp_ratio,
                    out_dim=out_dim,
                    slice_nums=slice_nums,
                    last_layer=(_ == n_layers - 1),
                )
                for _ in range(n_layers)
            ]
        )
        self.initialize_weights()
        self.placeholder = nn.Parameter(
            (1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float)
        )

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_grid(self, my_pos):
        # my_pos 1 N 3
        batchsize = my_pos.shape[0]

        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat(
            [batchsize, 1, self.ref, self.ref, 1]
        )
        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)
        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat(
            [batchsize, self.ref, 1, self.ref, 1]
        )
        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)
        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat(
            [batchsize, self.ref, self.ref, 1, 1]
        )
        grid_ref = (
            torch.cat((gridx, gridy, gridz), dim=-1)
            .cuda()
            .reshape(batchsize, self.ref**3, 3)
        )  # B 4 4 4 3

        pos = (
            torch.sqrt(
                torch.sum(
                    (my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2, dim=-1
                )
            )
            .reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref)
            .contiguous()
        )
        return pos

    def forward(self, data):
        # cfd_data, geom_data = data
        # x, fx, T = cfd_data.x, None, None
        # x  : node features
        # pos: node pos
        x, pos, fx = *data, None

        x = x[None, :, :]
        if self.unified_pos:
            new_pos = self.get_grid(pos[None, :, :])
            x = torch.cat((x, new_pos), dim=-1)

        if fx is not None:
            fx = torch.cat((x, fx), -1)
            fx = self.preprocess(fx)
        else:
            fx = self.preprocess(x)
            fx = fx + self.placeholder[None, None, :]

        for block in self.blocks:
            fx = block(fx)

        return fx[0]


class Transolver(nn.Module):
    def __init__(self, params) -> None:
        super().__init__()

        self.params = params
        self.nn = Model(
            space_dim=6,
            n_hidden=256,
            n_layers=8,
            fun_dim=0,
            n_head=8,
            mlp_ratio=2,
            out_dim=1,
            slice_num=32,
            unified_pos=0,
        )

    def forward(self, batch_data):
        # B C D H W
        batch_size = batch_data["node|pos"].shape[0]
        assert batch_size == 1
        pos = batch_data["node|pos"].reshape(-1, 3)
        node_features = batch_data["node|unit_norm_v"].reshape(-1, 3)
        x = torch.cat([pos, node_features], dim=1)
        output = self.nn(x, None)

        return output


class OriginTransolver(nn.Module):
    def __init__(
        self,
        message_passing_num,
        node_input_size,
        edge_input_size,
        cell_input_size,
        node_output_size,
        edge_output_size,
        cell_output_size,
        hidden_size,
        normlizer_steps,
        device,
        model_dir=None,
    ) -> None:
        super().__init__()
        self._device = device
        self.node_input_size = node_input_size
        self.edge_input_size = edge_input_size
        self.cell_input_size = cell_input_size
        self.model_dir = model_dir

        self.model = Model(
            space_dim=6,
            n_layers=message_passing_num,
            n_hidden=hidden_size,
            dropout=0,
            n_head=8,
            act="gelu",
            mlp_ratio=2,
            fun_dim=0,
            out_dim=1,
            slice_num=32,
            ref=8,
            unified_pos=False,
        ).to(device)

        self.node_normlizer = Normalizer(
            size=node_input_size,
            max_accumulations=normlizer_steps,
            epsilon=1e-8,
            device=device,
        )

        self.initialize_weights()

        print("Simulator model initialized")

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def update_node_attr(
        self,
        node_attr,
        is_training=False,
    ):

        norm_node_feature = self.node_normlizer(node_attr, accumulate=is_training)

        return norm_node_feature

    def forward(
        self,
        graph_cell=None,
        graph_edge=None,
        graph_node=None,
        params=None,
        is_training=True,
    ):
        """************************FORWARD**********************"""
        # forward model
        graph_node.x = self.update_node_attr(
            node_attr=graph_node.x,
            is_training=is_training,
        )

        predicted_node_uvp = self.model(
            x=graph_node.x, graph_node=graph_node, params=params
        )

        return predicted_node_uvp

    def load_checkpoint(
        self, optimizer=None, scheduler=None, ckpdir=None, device=None, is_training=True
    ):
        if ckpdir is None:
            ckpdir = self.model_dir
        dicts = torch.load(ckpdir, map_location=device)
        self.load_state_dict(dicts["model"])
        keys = list(dicts.keys())
        keys.remove("model")
        if optimizer is not None:
            if type(optimizer) is not list:
                optimizer = [optimizer]
            for i, o in enumerate(optimizer):
                o.load_state_dict(dicts["optimizer{}".format(i)])
                keys.remove("optimizer{}".format(i))

        if scheduler is not None:
            if type(scheduler) is not list:
                scheduler = [scheduler]
            for i, s in enumerate(scheduler):
                s.load_state_dict(dicts["scheduler{}".format(i)])
                keys.remove("scheduler{}".format(i))

        if not is_training:
            for key in keys.copy():
                if key.find("optimizer") >= 0:
                    keys.remove(key)
                elif key.find("scheduler") >= 0:
                    keys.remove(key)

        print("Simulator model and optimizer/scheduler loaded checkpoint %s" % ckpdir)

    def save_checkpoint(self, path=None, optimizer=None, scheduler=None):
        if path is None:
            path = self.model_dir

        model = self.state_dict()

        to_save = {"model": model}

        if type(optimizer) is not list:
            optimizer = [optimizer]
        for i, o in enumerate(optimizer):
            to_save.update({"optimizer{}".format(i): o.state_dict()})

        if type(scheduler) is not list:
            scheduler = [scheduler]
        for i, s in enumerate(scheduler):
            # to_save.update({"scheduler{}".format(i): s.get_variable()})
            to_save.update({"scheduler{}".format(i): s.state_dict()})

        torch.save(to_save, path)
        print("Simulator model saved at %s" % path)


class PretrainTransolver(nn.Module):
    def __init__(
        self,
        message_passing_num,
        node_input_size,
        edge_input_size,
        cell_input_size,
        node_output_size,
        edge_output_size,
        cell_output_size,
        hidden_size,
        normlizer_steps,
        device,
        pretrain=True,
        model_dir=None,
    ) -> None:
        super().__init__()
        self._device = device
        self.node_input_size = node_input_size
        self.edge_input_size = edge_input_size
        self.cell_input_size = cell_input_size
        self.model_dir = model_dir

        self.recon_pos_model = Model(
            space_dim=6,
            n_layers=message_passing_num // 2,
            n_hidden=hidden_size,
            dropout=0,
            n_head=8,
            act="gelu",
            mlp_ratio=2,
            fun_dim=0,
            out_dim=3,
            slice_num=32,
            ref=8,
            unified_pos=False,
        ).to(device)

        self.regressor = Model(
            space_dim=6,
            n_layers=message_passing_num // 2,
            n_hidden=hidden_size,
            dropout=0,
            n_head=8,
            act="gelu",
            mlp_ratio=2,
            fun_dim=0,
            out_dim=1,
            slice_num=32,
            ref=8,
            unified_pos=False,
        ).to(device)

        self.node_normlizer = Normalizer(
            size=node_input_size,
            max_accumulations=normlizer_steps,
            epsilon=1e-8,
            device=device,
        )

        self.initialize_weights()

        print("Simulator model initialized")

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def update_node_attr(
        self,
        node_attr,
        is_training=False,
    ):

        norm_node_feature = self.node_normlizer(node_attr, accumulate=is_training)

        return norm_node_feature

    def forward(
        self,
        graph_cell=None,
        graph_edge=None,
        graph_node=None,
        params=None,
        pretrain=False,
        is_training=True,
    ):
        """************************FORWARD**********************"""
        # forward model
        graph_node.x = self.update_node_attr(
            node_attr=graph_node.x,
            is_training=is_training,
        )

        if pretrain:
            self.regressor.eval()
            self.recon_pos_model.train()

            # # 以50%概率掩盖30%的点
            # if torch.rand(1).item() < 0.5:
            mask = torch.rand(graph_node.x.size(0), device=graph_node.x.device) < 0.3
            graph_node.x[mask] = 0.0

            predicted_node_pos = self.recon_pos_model(
                x=graph_node.x, graph_node=graph_node, params=params
            )

            return predicted_node_pos

        else:
            self.recon_pos_model.eval()
            self.regressor.train()

            with torch.no_grad():

                fx = self.recon_pos_model.encode(graph_node.x, graph_node.pos)

                prefx = self.recon_pos_model.process(fx)
                
            enc = self.regressor.encode(graph_node.x, graph_node.pos)
            
            fx = self.regressor.process(enc+prefx)

            predicted_node_uvp = self.regressor.decode(fx, graph_node.num_graphs)

            return predicted_node_uvp

    def load_checkpoint(
        self, optimizer=None, scheduler=None, ckpdir=None, device=None, is_training=True
    ):
        if ckpdir is None:
            ckpdir = self.model_dir
        dicts = torch.load(ckpdir, map_location=device)
        self.load_state_dict(dicts["model"])
        keys = list(dicts.keys())
        keys.remove("model")
        if optimizer is not None:
            if type(optimizer) is not list:
                optimizer = [optimizer]
            for i, o in enumerate(optimizer):
                o.load_state_dict(dicts["optimizer{}".format(i)])
                keys.remove("optimizer{}".format(i))

        if scheduler is not None:
            if type(scheduler) is not list:
                scheduler = [scheduler]
            for i, s in enumerate(scheduler):
                s.load_state_dict(dicts["scheduler{}".format(i)])
                keys.remove("scheduler{}".format(i))

        if not is_training:
            for key in keys.copy():
                if key.find("optimizer") >= 0:
                    keys.remove(key)
                elif key.find("scheduler") >= 0:
                    keys.remove(key)

        print("Simulator model and optimizer/scheduler loaded checkpoint %s" % ckpdir)

    def save_checkpoint(self, path=None, optimizer=None, scheduler=None):
        if path is None:
            path = self.model_dir

        model = self.state_dict()

        to_save = {"model": model}

        if type(optimizer) is not list:
            optimizer = [optimizer]
        for i, o in enumerate(optimizer):
            to_save.update({"optimizer{}".format(i): o.state_dict()})

        if type(scheduler) is not list:
            scheduler = [scheduler]
        for i, s in enumerate(scheduler):
            # to_save.update({"scheduler{}".format(i): s.get_variable()})
            to_save.update({"scheduler{}".format(i): s.state_dict()})

        torch.save(to_save, path)
        print("Simulator model saved at %s" % path)
